# %% 讯飞chat模型
from langchain_community.chat_models import ChatSparkLLM

import os

from langchain_core.tools import tool

# 加载环境变量
# 获取key详见: https://blog.csdn.net/weixin_56649281/article/details/136569427
appid = os.getenv('iflytek_spark_app_id')
apikey = os.getenv('iflytek_spark_api_key')
apisecret = os.getenv('iflytek_spark_api_secret')
apiurl = os.getenv('iflytek_spark_api_url')  # wss://spark-api.xf-yun.com/v1.1/chat
apimodel = os.getenv('iflytek_spark_model')  # lite

# 讯飞chat模型
llm = ChatSparkLLM(
    request_timeout=180,
    spark_api_url=apiurl,
    spark_llm_domain=apimodel,
    spark_app_id=appid, spark_api_key=apikey, spark_api_secret=apisecret
)

# %% 使用 Hugging Face Hub 上的嵌入模型
# from langchain_community.embeddings import HuggingFaceEmbeddings
#
# embeddings = HuggingFaceEmbeddings(model_name="xlm-roberta-large",
#                                    cache_folder='./models/huggingface/')
from local_embedding import LocalFastAPIEmbeddings

embeddings = LocalFastAPIEmbeddings()
# %% 搜索 agent
from langchain_community.utilities.tavily_search import TavilySearchAPIWrapper
from langchain_community.tools.tavily_search import TavilySearchResults

# https://blog.csdn.net/lovechris00/article/details/143053341
tavily_api_key = os.getenv('tavily_api_key')
search = TavilySearchAPIWrapper(tavily_api_key=tavily_api_key)
tavily_tool = TavilySearchResults(api_wrapper=search)
search_agent = llm

import wikipedia
from typing import List, Dict


# @tool
# def wikipedia_search(input: str, max_results: int = 5) -> List[Dict[str, str]]:
#     """
#     使用维基百科搜索并返回多个相关页面的摘要和链接。
#     返回格式：
#     [
#         {"content": "摘要内容", "url": "https://en.wikipedia.org/wiki/..."},
#         ...
#     ]
#     """
#     try:
#         # 搜索关键词相关的页面标题
#         page_titles = wikipedia.search(input, results=max_results)
#
#         results = []
#         for title in page_titles:
#             try:
#                 page = wikipedia.page(title)
#                 # 获取wiki页面的摘要
#                 summary = wikipedia.summary(title, sentences=5)
#                 results.append({
#                     "content": summary,
#                     "url": page.url
#                 })
#             except wikipedia.exceptions.PageError:
#                 continue
#         return results
#     except Exception as e:
#         return [{"content": f"搜索出错: {str(e)}", "url": ""}]

from langchain_core.documents import Document  # 或者使用简单 dict 封装


@tool
def wikipedia_search(input: str, max_results: int = 5) -> Document:
    """
    使用维基百科搜索并返回结果，包装成具有 .content 属性的对象。
    """
    try:
        page_titles = wikipedia.search(input, results=max_results)
        results = []
        for title in page_titles:
            try:
                page = wikipedia.page(title)
                summary = wikipedia.summary(title, sentences=5)
                results.append(f"{summary} ({page.url})")
            except wikipedia.exceptions.PageError:
                continue
        content = "\n\n".join(results)
        return Document(page_content=content, metadata={})
    except Exception as e:
        return Document(page_content=f"搜索出错: {str(e)}", metadata={})


web_search_tool = wikipedia_search
